#pragma once

#include <string>
#include <vector>
#include <map>
#include <set>
#include <iostream>
#include <sstream>
#include <algorithm>
#include <stdexcept>
#include <cctype>

#include "os/path/path.hpp"
#include "argsparse/ArgsEnum.hpp"

namespace argparse {
    class ArgumentError : public std::runtime_error {
    public:
        explicit ArgumentError(const std::string& message) : std::runtime_error(message) {}
    };

    class ArgumentParser {
    private:
        struct Argument {
            std::string name;           // 参数长名称 (e.g., "--input")
            std::string short_name;      // 参数短名称 (e.g., "-i")
            std::string help;            // 帮助信息
            std::string default_value;   // 默认值
            bool required;               // 是否必需
            bool has_value;              // 是否有值
            std::string type;            // 参数类型
            std::vector<std::string> choices; // 可选值列表
            std::string dest;            // 目标字段名
            bool is_positional;          // 是否是位置参数
            
            Argument() : required(false), has_value(true), is_positional(false) {}
        };
        
        std::string prog_name;
        std::string description;
        std::string epilog;
        std::vector<Argument> arguments; // 使用vector存储参数，保持添加顺序
        std::map<std::string, size_t> option_map; // 选项名称到参数索引的映射
        std::map<std::string, std::string> parsed_args;
        std::vector<std::string> positional_arg_names;
        
        // 获取参数索引
        size_t get_argument_index(const std::string& name) const {
            auto it = option_map.find(name);
            if (it != option_map.end()) {
                return it->second;
            }
            throw ArgumentError("unrecognized argument: " + name);
        }
        
        // 获取参数引用
        Argument& get_argument(const std::string& name) {
            return arguments[get_argument_index(name)];
        }
        
        // 获取参数常量引用
        const Argument& get_argument(const std::string& name) const {
            return arguments[get_argument_index(name)];
        }
        
        // 添加选项映射
        void add_option_mapping(const std::string& name, size_t index) {
            if (!name.empty()) {
                if (option_map.find(name) != option_map.end()) {
                    throw std::runtime_error("argument already exists: " + name);
                }
                option_map[name] = index;
            }
        }
        
        // 生成目标字段名
        std::string generate_dest(const std::string& name) {
            if (name.substr(0, 2) == "--") {
                return name.substr(2);
            } else if (name.substr(0, 1) == "-") {
                return name.substr(1);
            }
            return name;
        }
        
    public:
        ArgumentParser(const std::string& prog = "", const std::string& desc = "", const std::string& epi = "")
            : prog_name(prog), description(desc), epilog(epi) {}
        
        // 添加位置参数
        ArgumentParser& add_argument(const std::string& name) {
            Argument arg;
            arg.name = name;
            arg.dest = generate_dest(name);
            arg.is_positional = true;
            
            arguments.push_back(arg);
            size_t index = arguments.size() - 1;
            add_option_mapping(name, index);
            positional_arg_names.push_back(name);
            
            return *this;
        }
        
        // 添加可选参数
        ArgumentParser& add_argument(const std::string& name, const std::string& short_name, const std::string& help = "") {
            Argument arg;
            arg.name = name;
            arg.short_name = short_name;
            arg.help = help;
            arg.dest = generate_dest(name);
            arg.is_positional = false;
            
            arguments.push_back(arg);
            size_t index = arguments.size() - 1;
            add_option_mapping(name, index);
            add_option_mapping(short_name, index);
            
            return *this;
        }
        
        // 设置参数属性
        ArgumentParser& help(const std::string& help_text) {
            if (!arguments.empty()) {
                arguments.back().help = help_text;
            }
            return *this;
        }
        
        ArgumentParser& required(bool req = true) {
            if (!arguments.empty()) {
                arguments.back().required = req;
            }
            return *this;
        }
        
        ArgumentParser& default_value(const std::string& default_val) {
            if (!arguments.empty()) {
                arguments.back().default_value = default_val;
                arguments.back().required = false;
            }
            return *this;
        }
        
        ArgumentParser& action(const StoreType &act) {
            if (!arguments.empty()) {
                Argument& arg = arguments.back();
                if (act == StoreTrue || act == StoreFalse) {
                    arg.has_value = false;
                    arg.default_value = (act == StoreFalse) ? "false" : "true";
                }
            }
            return *this;
        }
        
        ArgumentParser& type(const ArgsType& type) {
            if (!arguments.empty()) {
                arguments.back().type = type;
            }
            return *this;
        }
        
        ArgumentParser& choices(const std::vector<std::string>& choice_list) {
            if (!arguments.empty()) {
                arguments.back().choices = choice_list;
            }
            return *this;
        }
        
        // 解析参数
        std::map<std::string, std::string> parse_args(int argc, char* argv[]) {
            if (prog_name.empty() && argc > 0) {
                prog_name = os::path::basename(argv[0]);
            }
            
            std::vector<std::string> args;
            for (int i = 1; i < argc; ++i) {
                args.push_back(argv[i]);
            }
            
            return parse_args(args);
        }
        
        std::map<std::string, std::string> parse_args(const std::vector<std::string>& args) {
            parsed_args.clear();
            
            // 设置默认值
            for (const auto& arg : arguments) {
                if (!arg.default_value.empty()) {
                    parsed_args[arg.dest] = arg.default_value;
                } else if (!arg.has_value) {
                    parsed_args[arg.dest] = "false";
                }
            }
            
            std::set<std::string> seen_dests;
            std::vector<std::string> positional_values;
            
            for (size_t i = 0; i < args.size(); ++i) {
                const std::string& arg_str = args[i];
                
                if (arg_str.substr(0, 2) == "--" || (arg_str.substr(0, 1) == "-" && arg_str.length() > 1)) {
                    // 处理可选参数
                    const Argument& argument = get_argument(arg_str);
                    seen_dests.insert(argument.dest);
                    
                    if (!argument.has_value) {
                        // store_true/store_false 类型
                        parsed_args[argument.dest] = (argument.default_value == "false") ? "true" : "false";
                    } else {
                        // 需要值的参数
                        if (i + 1 >= args.size()) {
                            throw ArgumentError("argument " + arg_str + ": expected one argument");
                        }
                        
                        std::string value = args[++i];
                        
                        // 检查选择列表
                        if (!argument.choices.empty()) {
                            if (std::find(argument.choices.begin(), argument.choices.end(), value) == argument.choices.end()) {
                                std::string choices_str;
                                for (size_t j = 0; j < argument.choices.size(); ++j) {
                                    if (j > 0) choices_str += ", ";
                                    choices_str += "'" + argument.choices[j] + "'";
                                }
                                throw ArgumentError("argument " + arg_str + ": invalid choice: '" + value + "' (choose from " + choices_str + ")");
                            }
                        }
                        
                        parsed_args[argument.dest] = value;
                    }
                } else {
                    // 位置参数
                    positional_values.push_back(arg_str);
                }
            }
            
            // 检查位置参数数量
            if (positional_values.size() > positional_arg_names.size()) {
                throw ArgumentError("too many positional arguments");
            }
            
            // 分配位置参数值
            for (size_t i = 0; i < positional_values.size(); ++i) {
                const std::string& name = positional_arg_names[i];
                const Argument& argument = get_argument(name);
                parsed_args[argument.dest] = positional_values[i];
                seen_dests.insert(argument.dest);
            }
            
            // 检查必需参数
            for (const auto& arg : arguments) {
                if (arg.required && seen_dests.find(arg.dest) == seen_dests.end()) {
                    throw ArgumentError("the following arguments are required: " + arg.name);
                }
            }
            
            return parsed_args;
        }
        
        // 获取参数值
        std::string get(const std::string& name) const {
            // 首先尝试直接查找dest
            auto it = parsed_args.find(name);
            if (it != parsed_args.end()) {
                return it->second;
            }
            
            // 尝试查找参数定义
            try {
                const Argument& arg = get_argument(name);
                auto it = parsed_args.find(arg.dest);
                if (it != parsed_args.end()) {
                    return it->second;
                }
                
                // 如果有默认值，返回默认值
                if (!arg.default_value.empty()) {
                    return arg.default_value;
                }
                
                // 如果是store_true/store_false类型，返回默认值
                if (!arg.has_value) {
                    return arg.default_value.empty() ? "false" : arg.default_value;
                }
                
                // 参数未提供且没有默认值
                throw ArgumentError("argument not provided: " + name);
            } catch (const ArgumentError&) {
                // 未找到参数定义
                throw ArgumentError("unknown argument: " + name);
            }
        }
        
        // 获取整数参数值
        int get_int(const std::string& name) const {
            std::string value = get(name);
            try {
                return std::stoi(value);
            } catch (const std::exception& e) {
                throw ArgumentError("invalid integer value for argument " + name + ": " + value);
            }
        }
        
        // 获取浮点数参数值
        double get_float(const std::string& name) const {
            std::string value = get(name);
            try {
                return std::stod(value);
            } catch (const std::exception& e) {
                throw ArgumentError("invalid float value for argument " + name + ": " + value);
            }
        }
        
        // 获取布尔参数值
        bool get_bool(const std::string& name) const {
            std::string value = get(name);
            std::string lower_value;
            std::transform(value.begin(), value.end(), std::back_inserter(lower_value),
                           [](unsigned char c){ return std::tolower(c); });
            
            if (lower_value == "true" || lower_value == "1" || lower_value == "yes" || lower_value == "on") {
                return true;
            }
            if (lower_value == "false" || lower_value == "0" || lower_value == "no" || lower_value == "off") {
                return false;
            }
            throw ArgumentError("invalid boolean value for argument " + name + ": " + value);
        }
        
        // 打印帮助信息
        void print_help() const {
            std::cout << "usage: " << (prog_name.empty() ? "program" : prog_name);
            
            // 打印位置参数
            for (const auto& name : positional_arg_names) {
                const Argument& arg = get_argument(name);
                std::cout << " " << arg.name;
            }
            
            // 打印可选参数
            std::cout << " [options]" << std::endl;
            
            // 打印描述
            if (!description.empty()) {
                std::cout << std::endl << description << std::endl;
            }
            
            // 打印位置参数
            if (!positional_arg_names.empty()) {
                std::cout << std::endl << "positional arguments:" << std::endl;
                for (const auto& name : positional_arg_names) {
                    const Argument& arg = get_argument(name);
                    std::cout << "  " << arg.name;
                    if (!arg.help.empty()) {
                        std::cout << " " << arg.help;
                    }
                    if (!arg.default_value.empty()) {
                        std::cout << " (default: " << arg.default_value << ")";
                    }
                    std::cout << std::endl;
                }
            }
            
            // 打印可选参数
            bool has_optional = false;
            for (const auto& arg : arguments) {
                if (!arg.is_positional) {
                    if (!has_optional) {
                        std::cout << std::endl << "optional arguments:" << std::endl;
                        has_optional = true;
                    }
                    
                    std::cout << "  " << arg.name;
                    if (!arg.short_name.empty()) {
                        std::cout << ", " << arg.short_name;
                    }
                    
                    if (!arg.help.empty()) {
                        std::cout << " " << arg.help;
                    }
                    
                    if (!arg.default_value.empty()) {
                        std::cout << " (default: " << arg.default_value << ")";
                    }
                    
                    if (!arg.choices.empty()) {
                        std::cout << " (choices: ";
                        for (size_t i = 0; i < arg.choices.size(); ++i) {
                            if (i > 0) std::cout << ", ";
                            std::cout << arg.choices[i];
                        }
                        std::cout << ")";
                    }
                    
                    std::cout << std::endl;
                }
            }
            
            // 打印结尾信息
            if (!epilog.empty()) {
                std::cout << std::endl << epilog << std::endl;
            }
        }
    };
}